import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import time

def get_bin_reg_loss(pred_reg, reg_label, loc_scope, loc_bin_size,
                    azimuth_scope=360, azimuth_bin_size=60,
                    elevation_scope=75, elevation_bin_size=15,
                    width_scope=0.04, width_bin_size=0.01,
                    grasp_angle_scope=180, grasp_angle_bin_size=30):

    """
    Bin-based 3D bounding boxes regression loss. See https://arxiv.org/abs/1812.04244 for more details.
    
    :param pred_reg: (N, C)
    :param reg_label: (N, 3) [dx, dy, dz]
    :param loc_scope: constant
    :param loc_bin_size: constant
    :param num_head_bin: constant
    :param anchor_size: (N, 3) or (3)
    :param get_xz_fine:
    :param get_y_by_bin:
    :param loc_scope:
    :param loc_bin_size:
    :param get_ry_fine:
    :return:
    """
    per_loc_bin_num = int(loc_scope / loc_bin_size) * 2

    reg_loss_dict = {}
    loc_loss = 0

    # xz localization loss
    x_offset_label, y_offset_label, z_offset_label = reg_label[:, 0], reg_label[:, 1], reg_label[:, 2]
    azimuth_offset_label, elevation_offset_label = reg_label[:, 3], reg_label[:, 4]
    width_offset_label = reg_label[:, 6]
    grasp_angle_offset_label = reg_label[:, 7]
    
    x_shift = torch.clamp(x_offset_label + loc_scope, 0, loc_scope * 2 - 1e-4)
    z_shift = torch.clamp(z_offset_label + loc_scope, 0, loc_scope * 2 - 1e-4)
    x_bin_label = (x_shift / loc_bin_size).floor().long()
    z_bin_label = (z_shift / loc_bin_size).floor().long()

    x_bin_l, x_bin_r = 0, per_loc_bin_num
    z_bin_l, z_bin_r = per_loc_bin_num, per_loc_bin_num * 2
    start_offset = z_bin_r

    loss_x_bin = F.cross_entropy(pred_reg[:, x_bin_l: x_bin_r], x_bin_label)
    loss_z_bin = F.cross_entropy(pred_reg[:, z_bin_l: z_bin_r], z_bin_label)

    reg_loss_dict['loss_x_bin'] = loss_x_bin.item()
    reg_loss_dict['loss_z_bin'] = loss_z_bin.item()
    loc_loss += loss_x_bin + loss_z_bin
    # get x, y, z localization loss
    x_res_l, x_res_r = per_loc_bin_num * 2, per_loc_bin_num * 3
    z_res_l, z_res_r = per_loc_bin_num * 3, per_loc_bin_num * 4
    start_offset = z_res_r

    # modified by weiwei 2020.10.28
    #x_res_label = x_shift - (x_bin_label.float() * loc_bin_size + loc_bin_size / 2)
    #z_res_label = z_shift - (z_bin_label.float() * loc_bin_size + loc_bin_size / 2)
    x_res_label = x_shift - (x_bin_label.float() * loc_bin_size)
    z_res_label = z_shift - (z_bin_label.float() * loc_bin_size)
    x_res_norm_label = x_res_label / loc_bin_size
    z_res_norm_label = z_res_label / loc_bin_size

    x_bin_onehot = torch.cuda.FloatTensor(x_bin_label.size(0), per_loc_bin_num).zero_()
    x_bin_onehot.scatter_(1, x_bin_label.view(-1, 1).long(), 1)
    z_bin_onehot = torch.cuda.FloatTensor(z_bin_label.size(0), per_loc_bin_num).zero_()
    z_bin_onehot.scatter_(1, z_bin_label.view(-1, 1).long(), 1)

    loss_x_res = F.smooth_l1_loss((torch.sigmoid(pred_reg[:, x_res_l: x_res_r]) * x_bin_onehot).sum(dim=1), x_res_norm_label)
    loss_z_res = F.smooth_l1_loss((torch.sigmoid(pred_reg[:, z_res_l: z_res_r]) * z_bin_onehot).sum(dim=1), z_res_norm_label)
    reg_loss_dict['loss_x_res'] = loss_x_res.item()
    reg_loss_dict['loss_z_res'] = loss_z_res.item()
    loc_loss += loss_x_res + loss_z_res

    # y localization loss
    y_bin_l, y_bin_r = start_offset, start_offset + per_loc_bin_num
    y_res_l, y_res_r = y_bin_r, y_bin_r + per_loc_bin_num
    start_offset = y_res_r

    y_shift = torch.clamp(y_offset_label + loc_scope, 0, loc_scope * 2 - 1e-4)
    y_bin_label = (y_shift / loc_bin_size).floor().long()
    #y_res_label = y_shift - (y_bin_label.float() * loc_bin_size + loc_bin_size / 2)
    y_res_label = y_shift - (y_bin_label.float() * loc_bin_size)
    y_res_norm_label = y_res_label / loc_bin_size

    y_bin_onehot = torch.cuda.FloatTensor(y_bin_label.size(0), per_loc_bin_num).zero_()
    y_bin_onehot.scatter_(1, y_bin_label.view(-1, 1).long(), 1)

    loss_y_bin = F.cross_entropy(pred_reg[:, y_bin_l: y_bin_r], y_bin_label)
    loss_y_res = F.smooth_l1_loss((torch.sigmoid(pred_reg[:, y_res_l: y_res_r]) * y_bin_onehot).sum(dim=1), y_res_norm_label)

    reg_loss_dict['loss_y_bin'] = loss_y_bin.item()
    reg_loss_dict['loss_y_res'] = loss_y_res.item()

    loc_loss += loss_y_bin + loss_y_res
    # size loss

    #approach_azimuth_scope=360, approach_azimuth_bin_size=30,
    #approach_elevation_scope=75, approach_elevation_bin_size=15
    per_azimuth_bin_num = int(azimuth_scope / azimuth_bin_size)
    per_elevation_bin_num = int(elevation_scope / elevation_bin_size)

    azimuth_bin_l, azimuth_bin_r = start_offset, start_offset + per_azimuth_bin_num
    azimuth_res_l, azimuth_res_r = azimuth_bin_r, azimuth_bin_r + per_azimuth_bin_num
    start_offset = azimuth_res_r

    azimuth_shift = torch.clamp(azimuth_offset_label, 0, azimuth_scope - 1e-4)
    azimuth_bin_label = (azimuth_shift / azimuth_bin_size).floor().long()
    #azimuth_res_label = azimuth_shift - (azimuth_bin_label.float() * azimuth_bin_size + azimuth_bin_size / 2)
    azimuth_res_label = azimuth_shift - (azimuth_bin_label.float() * azimuth_bin_size)
    azimuth_res_norm_label = azimuth_res_label / azimuth_bin_size

    azimuth_bin_onehot = torch.cuda.FloatTensor(azimuth_bin_label.size(0), per_azimuth_bin_num).zero_()
    azimuth_bin_onehot.scatter_(1, azimuth_bin_label.view(-1, 1).long(), 1)

    loss_azimuth_bin = F.cross_entropy(pred_reg[:, azimuth_bin_l: azimuth_bin_r], azimuth_bin_label)
    loss_azimuth_res = F.smooth_l1_loss((torch.sigmoid(pred_reg[:, azimuth_res_l: azimuth_res_r]) * azimuth_bin_onehot).sum(dim=1), azimuth_res_norm_label)
    reg_loss_dict['loss_azimuth_bin'] = loss_azimuth_bin.item()
    reg_loss_dict['loss_azimuth_res'] = loss_azimuth_res.item()
    loc_loss += loss_azimuth_bin + loss_azimuth_res


    elevation_bin_l, elevation_bin_r = start_offset, start_offset + per_elevation_bin_num
    elevation_res_l, elevation_res_r = elevation_bin_r, elevation_bin_r + per_elevation_bin_num
    start_offset = elevation_res_r

    elevation_shift = torch.clamp(elevation_offset_label, 0, elevation_scope - 1e-4)
    elevation_bin_label = (elevation_shift / elevation_bin_size).floor().long()
    #elevation_res_label = elevation_shift - (elevation_bin_label.float() * elevation_bin_size + elevation_bin_size / 2)
    elevation_res_label = elevation_shift - (elevation_bin_label.float() * elevation_bin_size)
    elevation_res_norm_label = elevation_res_label / elevation_bin_size

    elevation_bin_onehot = torch.cuda.FloatTensor(elevation_bin_label.size(0), per_elevation_bin_num).zero_()
    elevation_bin_onehot.scatter_(1, elevation_bin_label.view(-1, 1).long(), 1)

    loss_elevation_bin = F.cross_entropy(pred_reg[:, elevation_bin_l: elevation_bin_r], elevation_bin_label)
    loss_elevation_res = F.smooth_l1_loss((torch.sigmoid(pred_reg[:, elevation_res_l: elevation_res_r]) * elevation_bin_onehot).sum(dim=1), elevation_res_norm_label)
    reg_loss_dict['loss_elevation_bin'] = loss_elevation_bin.item()
    reg_loss_dict['loss_elevation_res'] = loss_elevation_res.item()

    loc_loss += loss_elevation_bin + loss_elevation_res

    per_width_bin_num = int(width_scope / width_bin_size)
    width_bin_l, width_bin_r = start_offset, start_offset + per_width_bin_num
    width_res_l, width_res_r = width_bin_r, width_bin_r + per_width_bin_num
    start_offset = width_res_r

    width_shift = torch.clamp(width_offset_label, 0, width_scope - 1e-4)
    width_bin_label = (width_shift / width_bin_size).floor().long()
    #width_res_label = width_shift - (width_bin_label.float() * width_bin_size + width_bin_size / 2)
    width_res_label = width_shift - (width_bin_label.float() * width_bin_size)
    width_res_norm_label = width_res_label / width_bin_size

    width_bin_onehot = torch.cuda.FloatTensor(width_bin_label.size(0), per_width_bin_num).zero_()
    width_bin_onehot.scatter_(1, width_bin_label.view(-1, 1).long(), 1)

    loss_width_bin = F.cross_entropy(pred_reg[:, width_bin_l: width_bin_r], width_bin_label)
    loss_width_res = F.smooth_l1_loss((torch.sigmoid(pred_reg[:, width_res_l: width_res_r]) * width_bin_onehot).sum(dim=1), width_res_norm_label)
    reg_loss_dict['loss_width_bin'] = loss_width_bin.item()
    reg_loss_dict['loss_width_res'] = loss_width_res.item()
    loc_loss += loss_width_bin + loss_width_res

    per_grasp_angle_bin_num = int(grasp_angle_scope / grasp_angle_bin_size)
    grasp_angle_bin_l, grasp_angle_bin_r = start_offset, start_offset + per_grasp_angle_bin_num
    grasp_angle_res_l, grasp_angle_res_r = grasp_angle_bin_r, grasp_angle_bin_r + per_grasp_angle_bin_num
    start_offset = grasp_angle_res_r

    grasp_angle_shift = torch.clamp(grasp_angle_offset_label, 0, grasp_angle_scope - 1e-4)
    grasp_angle_bin_label = (grasp_angle_shift / grasp_angle_bin_size).floor().long()
    #grasp_angle_res_label = grasp_angle_shift - (grasp_angle_bin_label.float() * grasp_angle_bin_size + grasp_angle_bin_size / 2)
    grasp_angle_res_label = grasp_angle_shift - (grasp_angle_bin_label.float() * grasp_angle_bin_size)
    grasp_angle_res_norm_label = grasp_angle_res_label / grasp_angle_bin_size

    grasp_angle_bin_onehot = torch.cuda.FloatTensor(grasp_angle_bin_label.size(0), per_grasp_angle_bin_num).zero_()
    grasp_angle_bin_onehot.scatter_(1, grasp_angle_bin_label.view(-1, 1).long(), 1)

    loss_grasp_angle_bin = F.cross_entropy(pred_reg[:, grasp_angle_bin_l: grasp_angle_bin_r], grasp_angle_bin_label)
    loss_grasp_angle_res = F.smooth_l1_loss((torch.sigmoid(pred_reg[:, grasp_angle_res_l: grasp_angle_res_r]) * grasp_angle_bin_onehot).sum(dim=1), grasp_angle_res_norm_label)
    reg_loss_dict['loss_grasp_angle_bin'] = loss_grasp_angle_bin.item()
    reg_loss_dict['loss_grasp_angle_res'] = loss_grasp_angle_res.item()
    loc_loss += loss_grasp_angle_bin + loss_grasp_angle_res
    #print(reg_label[0]) 
    #print(x_bin_label[0], x_res_norm_label[0])
    #print(y_bin_label[0], y_res_norm_label[0])
    #print(z_bin_label[0], z_res_norm_label[0])
    #print(width_bin_label[0], width_res_norm_label[0])
    #print(azimuth_bin_label[0], azimuth_res_norm_label[0])
    #print(elevation_bin_label[0], elevation_res_norm_label[0])
    #print(grasp_angle_bin_label[0], grasp_angle_res_norm_label[0])

    reg_loss_dict['loss_loc'] = loc_loss

    return loc_loss, reg_loss_dict

